import gensim
from gensim.models import Word2Vec
from gensim.corpora import Dictionary
import sys

def getText(filename):
    with open(filename,'r',encoding='utf-8') as rf:
        lines=rf.readlines()
        all_text=[]
        all_length=[]
        for i,line in enumerate(lines):
            if i%2!=0:
                text=line.strip('\n')[6:len(line)-8]
                sp=text.split()
                len1=len(sp)
                all_text.append(sp)
                all_length.append(len1)
            if i==40000:
                break
    return all_text,all_length

# with open('a.txt','r',encoding='utf-8') as rf:
#     lines=rf.readlines()
#     all_text=[]
#     all_length=[]
#     for i,line in enumerate(lines):
#         if i%2!=0:
#             text=line.strip('\n')[6:len(line)-8]
#             sp=text.split()
#     #         len1=len(sp)
#             all_text.append(sp)
    #         all_length.append(len1)
text1,len1=getText('a.txt')
text2,len2=getText('b.txt')
# print(text1[0])
# sys.exit(0)

a_text=[]
for i in range(len(text1)):
    shuffle_text=[]
    if len1[i]>len2[i]:
        mol=int(len1[i]/len2[i])
        elen=len2[i]
        buf=len1[i]%len2[i]
        for m in range(elen):
            for j in range(mol):
                shuffle_text.append(text1[i].pop(0))
            shuffle_text.append(text2[i].pop(0))
        shuffle_text.extend(text1[i])
    else:
        mol=int(len2[i]/len1[i])
        elen=len1[i]
        buf=len2[i]%len1[i]
        for m in range(elen):
            for j in range(mol):
                shuffle_text.append(text2[i].pop(0))
            shuffle_text.append(text1[i].pop(0))
        shuffle_text.extend(text2[i])
    a_text.append(shuffle_text)


model = Word2Vec(a_text, size=300, window=16, min_count=5, workers=4,negative=15)
dct = Dictionary(a_text)
dct.save('word.dict')
model.save('word2vec.model')






